package db

import (
	"data-system/models"
	"strings"
)

type StageDataSet struct {
	NodeType string
	IdData   map[int64]map[string]interface{}
}

var Stage map[string][]*StageDataSet

func addStageData(stageKey string, stageDataSets []*StageDataSet) error {
	dataSets := Stage[stageKey]
	if dataSets == nil {
		Stage[stageKey] = stageDataSets
		return nil
	}
	for _, stageDataSet := range stageDataSets {
		var oldStageDataSet *StageDataSet
		for _, dataSet := range dataSets {
			if strings.EqualFold(dataSet.NodeType, stageDataSet.NodeType) {
				oldStageDataSet = dataSet
				break
			}
		}
		if oldStageDataSet == nil {
			oldStageDataSet = stageDataSet
		} else {
			oldIdData := oldStageDataSet.IdData
			for id, data := range stageDataSet.IdData {
				oldData := oldIdData[id]
				for key, value := range data {
					oldData[key] = value
				}
				oldIdData[id] = oldData
			}
			oldStageDataSet.IdData = oldIdData
		}
		dataSets = append(dataSets, stageDataSet)
	}
	return nil
}

func getStageData(stageKey string) (map[string][]models.NodeData, error) {
	var dataSetMap = make(map[string][]models.NodeData, 0)
	for key, dataSets := range Stage {
		if strings.EqualFold(key, stageKey) {
			var ds = make([]models.NodeData, 0)
			for _, dataSet := range dataSets {
				for id, data := range dataSet.IdData {
					ds = append(ds, models.NodeData{
						NodeType: dataSet.NodeType,
						Id:       id,
						Data:     data,
					})
				}
				dataSetMap[dataSet.NodeType] = ds
			}
		}
	}
	return dataSetMap, nil
}
